{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "65548359",
   "metadata": {},
   "source": [
    "# memory update details\n",
    "\n",
    "see [studio/directive_memory_bot.py](studio/directive_memory_bot.py)\n",
    "\n",
    "- Hey, how are you? I’m Evgeny, I’m a software engineer and I need your help with my project. This is a TaskManager Java Spring Boot application. I do not get how to configure security there! Also, I often come back to these answers later, so please include comments inside any code you share.\n",
    "- Sorry if I sound a bit stressed — I’m not very experienced with Spring Security. If you could break things down step by step, that would really help."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b37e2d95",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pydantic import BaseModel, Field\n",
    "\n",
    "class InteractionInstruction(BaseModel):\n",
    "    instruction: str = Field(\n",
    "        description=(\n",
    "            \"An independent, atomic instruction that describes one specific way the assistant should adapt its interaction to the user. \"\n",
    "            \"Each instruction should focus on a single behavior or style — do not combine multiple preferences into one. \"\n",
    "            \"Instructions must be phrased as imperatives. \"\n",
    "            \"Examples: \"\n",
    "            \"'Explain technical topics using a step-by-step format.', \"\n",
    "            \"'Use concise code examples.', \"\n",
    "            \"'Avoid adding explanations unless explicitly requested.', \"\n",
    "            \"'Include inline comments in code when requested.', \"\n",
    "            \"'Ask if the user wants a deeper explanation after giving an answer.'\"\n",
    "        )\n",
    "    )\n",
    "\n",
    "class InteractionInstructionList(BaseModel):\n",
    "    instructions: list[InteractionInstruction] = Field(\n",
    "        description=\"A collection of interaction-relevant insights to guide how the assistant should respond to this user.\"\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9c70a00b",
   "metadata": {},
   "outputs": [],
   "source": [
    "from trustcall import create_extractor\n",
    "from langchain_openai import ChatOpenAI\n",
    "from langchain.schema import HumanMessage\n",
    "\n",
    "model = ChatOpenAI(model=\"gpt-4o-mini\")\n",
    "\n",
    "trustcall_extractor = create_extractor(\n",
    "    model,\n",
    "    tools=[InteractionInstruction],\n",
    "    tool_choice=\"InteractionInstruction\",\n",
    "    enable_inserts=True\n",
    ")\n",
    "\n",
    "result = trustcall_extractor.invoke([HumanMessage(content=\"\"\"\n",
    "Hey, how are you? I am Evgeny, I am a junior software engineer and I need your help with \n",
    "my project. This is a TaskManager Java Spring Boot application. I do not get how\n",
    "to configure security there!\n",
    "Also, I often come back to these answers later, so please include comments inside any code you share.\n",
    "\"\"\")])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8cdb4e81",
   "metadata": {},
   "outputs": [],
   "source": [
    "for m in result[\"messages\"]:\n",
    "    m.pretty_print()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cde08e16",
   "metadata": {},
   "outputs": [],
   "source": [
    "existing_instructions = [InteractionInstruction(instruction=result[\"responses\"][0].instruction)]\n",
    "existing_instructions"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8b9b0f6e",
   "metadata": {},
   "source": [
    "## Update & Listen"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "92567109",
   "metadata": {},
   "outputs": [],
   "source": [
    "class ToolsListener:\n",
    "    def __init__(self):\n",
    "        self.tools = []\n",
    "\n",
    "    def __call__(self, execution):\n",
    "        runs = [execution]\n",
    "        while runs:\n",
    "            run = runs.pop()\n",
    "            if run.child_runs:\n",
    "                runs.extend(run.child_runs)\n",
    "            if run.run_type == \"chat_model\":\n",
    "                self.tools.append(\n",
    "                    run.outputs[\"generations\"][0][0][\"message\"][\"kwargs\"][\"tool_calls\"]\n",
    "                )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7b5e9c6c",
   "metadata": {},
   "outputs": [],
   "source": [
    "# prepare existing facts\n",
    "schema = \"InteractionInstruction\"\n",
    "\n",
    "memories = [(str(i + 1), schema, instruction.model_dump()) for i, instruction in enumerate(existing_instructions)]\n",
    "memories"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e3c42cac",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_core.messages import HumanMessage, AIMessage\n",
    "\n",
    "listener = ToolsListener()\n",
    "extractor_with_listener = trustcall_extractor.with_listeners(on_end=listener)\n",
    "\n",
    "\n",
    "conversation = [\n",
    "    HumanMessage(content=\"\"\"\n",
    "    Sorry if I sound a bit stressed — I’m not very experienced with Spring Security.\n",
    "    If you could break things down step by step, that would really help.\n",
    "    \"\"\"),\n",
    "\n",
    "    AIMessage(content=\"\"\"\n",
    "    Totally understood! No worries at all — I’ll walk you through it one step at a time.\n",
    "    Let’s start with the basic configuration: you’ll need to add the Spring Security dependency first. Here's how...\n",
    "    \"\"\"), \n",
    "\n",
    "    HumanMessage(content=\"\"\"\n",
    "    Thanks, but could you skip the long explanation? Just tell me exactly what I need to do.\n",
    "    \"\"\"),\n",
    "]\n",
    "\n",
    "# Update the instruction\n",
    "system_msg = \"\"\"\n",
    "Update existing directives about the user and create new ones based on the following conversation, \n",
    "ensuring each directive reflects one specific communication preference or behavior the assistant should follow.\n",
    "\"\"\"\n",
    "\n",
    "\n",
    "result = extractor_with_listener.invoke({\n",
    "    \"messages\": [system_msg] + conversation, \n",
    "    \"existing\": memories\n",
    "})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8bad4f83",
   "metadata": {},
   "outputs": [],
   "source": [
    "for m in result[\"response_metadata\"]: \n",
    "    print(m)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "12a7968e",
   "metadata": {},
   "outputs": [],
   "source": [
    "for m in result[\"messages\"]:\n",
    "    m.pretty_print()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3afa92b6",
   "metadata": {},
   "outputs": [],
   "source": [
    "listener.tools"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "625aad83",
   "metadata": {},
   "outputs": [],
   "source": [
    "def summarize_listener_tools(tools_data):\n",
    "    if not tools_data:\n",
    "        print(\"No tool calls recorded.\")\n",
    "        return\n",
    "\n",
    "    for i, tool_calls in enumerate(tools_data):\n",
    "        print(f\"\\n=== Tool Calls for Document {i} ===\")\n",
    "        for call in tool_calls:\n",
    "            name = call.get(\"name\")\n",
    "            args = call.get(\"args\", {})\n",
    "            if name == \"PatchDoc\":\n",
    "                doc_id = args.get(\"json_doc_id\", \"Unknown\")\n",
    "                plan = args.get(\"planned_edits\", \"No plan provided.\")\n",
    "                patches = args.get(\"patches\", [])\n",
    "\n",
    "                print(f\"\\n🛠️ Document {doc_id} updated:\")\n",
    "                print(f\"📝 Plan: {plan}\")\n",
    "                for patch in patches:\n",
    "                    op = patch.get(\"op\")\n",
    "                    path = patch.get(\"path\")\n",
    "                    value = patch.get(\"value\")\n",
    "                    print(f\"🔧 Patch: {op.upper()} {path} → {value}\")\n",
    "\n",
    "            else:\n",
    "                print(f\"\\n➕ New memory created:\")\n",
    "                print(f\"📌 Tool: {name}\")\n",
    "                print(f\"📄 Content: {args}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e405251a",
   "metadata": {},
   "outputs": [],
   "source": [
    "summarize_listener_tools(listener.tools)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "venv",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.13.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
